{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/lokwq/TextBrewer/blob/add_note_examples/sqaudv1.1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rOvyG9zChDJu"
   },
   "source": [
    "This notebook shows how to fine-tune a model on squadv1.1 dataset and how to distill the model with TextBrewer.\n",
    "\n",
    "Detailed Docs can be find here: https://github.com/airaria/TextBrewer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "PqJEpBdjiMKQ"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "device='cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sgDADHlOiSw4"
   },
   "outputs": [],
   "source": [
    "!pip install transformers\n",
    "!pip install textbrewer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UWyCSU6u0ido"
   },
   "source": [
    "### Prepare dataset to train teacher model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "nYjP-ZDOv8zR"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import timeit\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "from tqdm import tqdm_notebook\n",
    "from transformers import AdamW,BertConfig,BertForQuestionAnswering,BertTokenizer,get_linear_schedule_with_warmup,squad_convert_examples_to_features\n",
    "\n",
    "from transformers.data.metrics.squad_metrics import compute_predictions_logits,squad_evaluate\n",
    "from transformers.data.processors.squad import SquadResult, SquadV1Processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cir5Dfr8ie0Z",
    "outputId": "536e8619-a4dc-413b-c8de-1cbc091ceda8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-07-08 04:22:17--  https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\n",
      "Resolving rajpurkar.github.io (rajpurkar.github.io)... 185.199.108.153, 185.199.109.153, 185.199.110.153, ...\n",
      "Connecting to rajpurkar.github.io (rajpurkar.github.io)|185.199.108.153|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 30288272 (29M) [application/json]\n",
      "Saving to: ‘train-v1.1.json’\n",
      "\n",
      "train-v1.1.json     100%[===================>]  28.88M  78.2MB/s    in 0.4s    \n",
      "\n",
      "2021-07-08 04:22:18 (78.2 MB/s) - ‘train-v1.1.json’ saved [30288272/30288272]\n",
      "\n",
      "--2021-07-08 04:22:18--  https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\n",
      "Reusing existing connection to rajpurkar.github.io:443.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 4854279 (4.6M) [application/json]\n",
      "Saving to: ‘dev-v1.1.json’\n",
      "\n",
      "dev-v1.1.json       100%[===================>]   4.63M  --.-KB/s    in 0.02s   \n",
      "\n",
      "2021-07-08 04:22:19 (215 MB/s) - ‘dev-v1.1.json’ saved [4854279/4854279]\n",
      "\n",
      "FINISHED --2021-07-08 04:22:19--\n",
      "Total wall clock time: 1.7s\n",
      "Downloaded: 2 files, 34M in 0.4s (85.7 MB/s)\n"
     ]
    }
   ],
   "source": [
    "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0NR7hLfqvyxZ",
    "outputId": "01fb674d-fb22-41e0-db98-8f44873bbb84"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 442/442 [00:33<00:00, 13.29it/s]\n"
     ]
    }
   ],
   "source": [
    "processor = SquadV1Processor()\n",
    "examples = processor.get_train_examples('/content/')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JOjXEHwQwBRZ",
    "outputId": "2ed2e42c-95cc-4c4b-a7b4-667c7642277c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertForQuestionAnswering: ['cls.predictions.decoder.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
      "- This IS expected if you are initializing BertForQuestionAnswering from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForQuestionAnswering were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertForQuestionAnswering(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(28996, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (qa_outputs): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-cased')\n",
    "model = BertForQuestionAnswering.from_pretrained('bert-base-cased')\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GEcJPauczrjK",
    "outputId": "49111455-999e-4455-faf7-ce6e88b0369e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "convert squad examples to features: 100%|██████████| 87599/87599 [11:27<00:00, 127.42it/s]\n",
      "add example index and unique id: 100%|██████████| 87599/87599 [00:00<00:00, 761102.36it/s]\n"
     ]
    }
   ],
   "source": [
    "features,train_dataset = squad_convert_examples_to_features(\n",
    "    examples=examples,\n",
    "    tokenizer=tokenizer,\n",
    "    max_seq_length=384,\n",
    "    doc_stride=128,\n",
    "    max_query_length=64,\n",
    "    is_training=True,\n",
    "    return_dataset=\"pt\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "xNJhv3Y_wJxU"
   },
   "outputs": [],
   "source": [
    "train_sampler = RandomSampler(train_dataset)\n",
    "train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sH8Ivk5E2Cuf",
    "outputId": "710fd6d0-1119-44b3-a86d-b07d17f33eb4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[ 101, 1332, 1108,  ...,    0,    0,    0],\n",
       "         [ 101, 1327, 2523,  ...,    0,    0,    0],\n",
       "         [ 101, 1130, 1184,  ...,    0,    0,    0],\n",
       "         ...,\n",
       "         [ 101, 1731, 1242,  ...,    0,    0,    0],\n",
       "         [ 101, 1327, 1110,  ...,    0,    0,    0],\n",
       "         [ 101, 1327, 1160,  ...,    0,    0,    0]]),\n",
       " tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0],\n",
       "         [1, 1, 1,  ..., 0, 0, 0]]),\n",
       " tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
       "         [0, 0, 0,  ..., 0, 0, 0],\n",
       "         [0, 0, 0,  ..., 0, 0, 0],\n",
       "         ...,\n",
       "         [0, 0, 0,  ..., 0, 0, 0],\n",
       "         [0, 0, 0,  ..., 0, 0, 0],\n",
       "         [0, 0, 0,  ..., 0, 0, 0]]),\n",
       " tensor([ 81,  99,  21, 146, 122, 130, 191,  77]),\n",
       " tensor([ 82, 100,  21, 147, 125, 130, 192,  80]),\n",
       " tensor([0, 0, 0, 0, 0, 0, 0, 0]),\n",
       " tensor([[0., 1., 1.,  ..., 1., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
       "         ...,\n",
       "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 1., 1., 1.]]),\n",
       " tensor([0., 0., 0., 0., 0., 0., 0., 0.])]"
      ]
     },
     "execution_count": 9,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(train_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rXkT9uN79JLi"
   },
   "outputs": [],
   "source": [
    "#Start training\n",
    "import tqdm\n",
    "epochs = 2\n",
    "t_total = len(train_dataloader) * epochs\n",
    "\n",
    "# Prepare optimizer and schedule \n",
    "optimizer = AdamW(model.parameters(), lr=3e-5, eps=1e-8)\n",
    "scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=t_total)\n",
    "\n",
    "\n",
    "print(\"  Num examples = \", len(train_dataset))\n",
    "print(\"  Total optimization steps = \", t_total)\n",
    "\n",
    "steps = 1\n",
    "tr_loss = 0.0\n",
    "model.zero_grad()\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print('Epoch:{}'.format(epoch+1))\n",
    "    epoch_iterator = tqdm.notebook.tqdm(train_dataloader, desc=\"Iteration\", disable=False)\n",
    "    for step, batch in enumerate(epoch_iterator):\n",
    "\n",
    "        model.train()\n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "\n",
    "        inputs = {\n",
    "            \"input_ids\": batch[0],\n",
    "            \"attention_mask\": batch[1],\n",
    "            \"token_type_ids\": batch[2],\n",
    "            \"start_positions\": batch[3],\n",
    "            \"end_positions\": batch[4],\n",
    "        }\n",
    "\n",
    "        outputs = model(**inputs)\n",
    "        loss = outputs.loss\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        tr_loss += loss.item()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "        optimizer.step()\n",
    "        scheduler.step()  \n",
    "        model.zero_grad()\n",
    "        steps += 1\n",
    "        \n",
    "        # Log \n",
    "        if steps % 500 == 0:\n",
    "            print('steps = {}, logging_loss = {}'.format(steps,tr_loss))\n",
    "\n",
    "print(\" steps = %s, average loss = %s\", steps, tr_loss / steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k-QLh-3IACr1"
   },
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), '/content/drive/MyDrive/squad_teacher_model.pt') #save the teacher model weights to distill"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SQHyopKl0qB4"
   },
   "source": [
    "### Start distillation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "a8AY5MZMBqpv"
   },
   "outputs": [],
   "source": [
    "import textbrewer\n",
    "from textbrewer import GeneralDistiller\n",
    "from textbrewer import TrainingConfig, DistillationConfig\n",
    "from transformers import BertForSequenceClassification, BertConfig, AdamW,BertTokenizer\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g0SWHH47jaz6"
   },
   "outputs": [],
   "source": [
    "from transformers import BertForQuestionAnswering,BertConfig\n",
    "\n",
    "\n",
    "bert_config = BertConfig.from_json_file('/content/drive/MyDrive/TextBrewer-master/examples/student_config/bert_base_cased_config/bert_config.json')\n",
    "bert_config_T3 = BertConfig.from_json_file('/content/drive/MyDrive/TextBrewer-master/examples/student_config/bert_base_cased_config/bert_config_L3.json')\n",
    "\n",
    "bert_config.output_hidden_states = True\n",
    "bert_config_T3.output_hidden_states = True\n",
    "\n",
    "\n",
    "teacher_model = BertForQuestionAnswering(bert_config) \n",
    "student_model = BertForQuestionAnswering(bert_config_T3) \n",
    "\n",
    "teacher_model.load_state_dict(torch.load('/content/drive/MyDrive/squad_teacher_model.pt'))\n",
    "\n",
    "teacher_model.to(device=device)\n",
    "student_model.to(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KD_vql_pje_J"
   },
   "outputs": [],
   "source": [
    "num_epochs = 20\n",
    "num_training_steps = len(train_dataloader) * num_epochs\n",
    "# Optimizer and learning rate scheduler\n",
    "optimizer = AdamW(student_model.parameters(), lr=1e-4)\n",
    "\n",
    "scheduler_class = get_linear_schedule_with_warmup\n",
    "# arguments dict except 'optimizer'\n",
    "scheduler_args = {'num_warmup_steps':int(0.1*num_training_steps), 'num_training_steps':num_training_steps}\n",
    "\n",
    "\n",
    "def simple_adaptor(batch, model_outputs):\n",
    "    return {'logits': (model_outputs.start_logits,model_outputs.end_logits), \n",
    "        'hidden': model_outputs.hidden_states, \n",
    "        'attention': model_outputs.attentions}\n",
    "\n",
    "distill_config = DistillationConfig(\n",
    "    temperature = 1,\n",
    "    intermediate_matches=[{\"layer_T\":[0,0],  \"layer_S\":[0,0], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
    "               {\"layer_T\":[4,4],  \"layer_S\":[1,1], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
    "               {\"layer_T\":[8,8],  \"layer_S\":[2,2], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
    "               {\"layer_T\":[12,12], \"layer_S\":[3,3], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1}])\n",
    "train_config = TrainingConfig()\n",
    "\n",
    "distiller = GeneralDistiller(\n",
    "    train_config=train_config, distill_config=distill_config,\n",
    "    model_T=teacher_model, model_S=student_model, \n",
    "    adaptor_T=simple_adaptor, adaptor_S=simple_adaptor)\n",
    "\n",
    "def batch_postprocessor(batch):\n",
    "  return {\"input_ids\": batch[0],\n",
    "       \"attention_mask\": batch[1],\n",
    "       \"token_type_ids\": batch[2],\n",
    "       \"start_positions\": batch[3],\n",
    "       \"end_positions\": batch[4]}\n",
    "\n",
    "with distiller:\n",
    "    distiller.train(optimizer, train_dataloader, num_epochs, scheduler_class=scheduler_class, scheduler_args = scheduler_args, batch_postprocessor=batch_postprocessor, callback=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5i0ucOPulw4T"
   },
   "outputs": [],
   "source": [
    "bert_config_T3 = BertConfig.from_json_file('/content/drive/MyDrive/data/bert_config/bert_config_L3.json')\n",
    "test_model = BertForQuestionAnswering(bert_config_T3)\n",
    "\n",
    "test_model.load_state_dict(torch.load('/content/saved_models/gs2813.pkl'))\n",
    "test_model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ExZU8DTXMAkl",
    "outputId": "a2926358-328a-40d0-aab0-cb5b41f7116e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 48/48 [00:06<00:00,  7.61it/s]\n",
      "convert squad examples to features: 100%|██████████| 10570/10570 [01:24<00:00, 125.69it/s]\n",
      "add example index and unique id: 100%|██████████| 10570/10570 [00:00<00:00, 685666.02it/s]\n"
     ]
    }
   ],
   "source": [
    "processor = SquadV1Processor()\n",
    "examples = processor.get_dev_examples('/content/')\n",
    "\n",
    "features,eval_dataset = squad_convert_examples_to_features(\n",
    "    examples=examples,\n",
    "    tokenizer=tokenizer,\n",
    "    max_seq_length=384,\n",
    "    doc_stride=128,\n",
    "    max_query_length=64,\n",
    "    is_training=False,\n",
    "    return_dataset=\"pt\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VZScjwwolz1V"
   },
   "outputs": [],
   "source": [
    "# Evaluate\n",
    "import tqdm\n",
    "from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union\n",
    "eval_sampler = SequentialSampler(eval_dataset)\n",
    "eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=8)\n",
    "\n",
    "print(\"  Num examples = \", len(eval_dataset))\n",
    "\n",
    "all_results = []\n",
    "\n",
    "def to_tuple(self) -> Tuple[Any]:\n",
    "    \"\"\"\n",
    "    Convert self to a tuple containing all the attributes/keys that are not ``None``.\n",
    "    \"\"\"\n",
    "    return tuple(self[k] for k in self.keys())\n",
    "\n",
    "for batch in tqdm.notebook.tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
    "    test_model.eval()\n",
    "    batch = tuple(t.to(device) for t in batch)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        inputs = {\n",
    "            \"input_ids\": batch[0],\n",
    "            \"attention_mask\": batch[1],\n",
    "            \"token_type_ids\": batch[2],\n",
    "        }\n",
    "\n",
    "        feature_indices = batch[3]\n",
    "\n",
    "        outputs = model(**inputs)\n",
    "\n",
    "    for i, feature_index in enumerate(feature_indices):\n",
    "        eval_feature = features[feature_index.item()]\n",
    "        unique_id = int(eval_feature.unique_id)\n",
    "        #output = [output[i].detach().cpu().tolist() for output in outputs]\n",
    "        #start_logits, end_logits = output\n",
    "        output = [(output[i]).detach().cpu().tolist() for output in outputs.to_tuple()]\n",
    "        start_logits, end_logits = output\n",
    "        result = SquadResult(unique_id, start_logits, end_logits)\n",
    "        all_results.append(result)\n",
    "\n",
    "\n",
    "predictions = compute_predictions_logits(\n",
    "    examples,\n",
    "    features,\n",
    "    all_results,\n",
    "    n_best_size=20,\n",
    "    max_answer_length=30,\n",
    "    do_lower_case=False,\n",
    "    output_prediction_file=\"predictions.json\",\n",
    "    output_nbest_file=\"nbest_predictions.json\",\n",
    "    output_null_log_odds_file=None,\n",
    "    verbose_logging=False,\n",
    "    version_2_with_negative=False,\n",
    "    null_score_diff_threshold=0.0,\n",
    "    tokenizer=tokenizer,)\n",
    "\n",
    "# Compute the F1 and exact scores.\n",
    "results = squad_evaluate(examples, predictions)\n",
    "\n",
    "print(\"Results: {}\".format(results))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyMBSv726UpUDQzxX2vrJ89Q",
   "collapsed_sections": [],
   "include_colab_link": true,
   "mount_file_id": "1xvWkBaH_COw0O5HB0ZDvS-Ki7Fu7mh4z",
   "name": "aqdistill.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
